#include "PreDepthProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "Game/Component/MultiModel.h"
#include "LocalScene.h"

void PreDepthProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mSceneEntity = static_cast<LocalEntity*>(mGameScene->GetEntityByName("scene"));

	CameraUbo cameraUbo;
	cameraUbo.proj = mGameScene->GetCamera().GetProj();
	cameraUbo.view = mGameScene->GetCamera().GetView();
	cameraUbo.cameraPos = mGameScene->GetCamera().GetPosition();

	MultiModel& model = mSceneEntity->GetComponent<MultiModel>();
	mObjectSize = model.mVertexGroups.size();

	ObjectUbo objectUbo;
	objectUbo.model = mSceneEntity->GetTransform().GetWorldMatrix();


	ENQUEUE_RENDER_COMMAND(BuildModel)([this, objectUbo, cameraUbo](VkCommandBuffer commandBuffer)
		{
			mEmptyImage = Picture2DTexture::Builder()
				.SetPath("Block.png")
				.SetAutoMip()
				.Build(this, "emptyImage");

			mEmptySampler = RHISampler::Builder()
				.SetMaxLod(mEmptyImage->GetMipLevels())
				.Build(this, "emptySampler");


			CameraUbo cameraUboData = cameraUbo;

			mCameraBuffer = CustomBuffer::Builder()
				.SetData(&cameraUboData)
				.SetSize(sizeof(CameraUbo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build(this, "cameraUbo");



			ObjectUbo objectUboData = objectUbo;

			mSceneEntity->mObjectBuffer = CustomBuffer::Builder()
				.SetData(&objectUboData)
				.SetSize(sizeof(ObjectUbo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			mSceneEntity->mMaterialBuffer = CustomBuffer::Builder()
				.SetData(&mSceneEntity->mMaterial)
				.SetSize(sizeof(MaterialUbo))
				.SetUsage(VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT)
				.SetMemoryProperty(VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
				.Build();

			if (mSceneEntity->mMaterial.hasAlbedoMap == 1)
			{
				mSceneEntity->mAlbedoImage = Picture2DTexture::Builder()
					.SetPath(mSceneEntity->mAlbedoPath)
					.SetAutoMip()
					.Build();

				mSceneEntity->mAlbedoSampler = RHISampler::Builder()
					.SetMaxLod(mSceneEntity->mAlbedoImage->GetMipLevels())
					.Build();
			}

			if (mSceneEntity->mMaterial.hasNormalMap == 1)
			{
				mSceneEntity->mNormalImage = Picture2DTexture::Builder()
					.SetPath(mSceneEntity->mNormalPath)
					.SetAutoMip()
					.Build();

				mSceneEntity->mNormalSampler = RHISampler::Builder()
					.SetMaxLod(mSceneEntity->mNormalImage->GetMipLevels())
					.Build();
			}

			
			MultiModel& model = mSceneEntity->GetComponent<MultiModel>();

			mSceneEntity->mVertexBuffers.resize(mObjectSize);
			mSceneEntity->mIndexBuffers.resize(mObjectSize);
			mSceneEntity->mDescriptorMasks.resize(mObjectSize);

			for (uint32 i = 0; i < mObjectSize; ++i)
			{
				std::vector<EntityVertex> tempVertexs;
				tempVertexs.resize(model.mVertexGroups[i].size());
				for (uint32 j = 0; j < model.mVertexGroups[i].size(); ++j)
				{
					tempVertexs[j].position = model.mVertexGroups[i][j].Position;
					tempVertexs[j].normal = model.mVertexGroups[i][j].Normal;
					tempVertexs[j].uv = model.mVertexGroups[i][j].TexC;
				}

				mSceneEntity->mVertexBuffers[i] = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(EntityVertex))
					.Build();

				mSceneEntity->mIndexBuffers[i] = IndexBuffer::Builder()
					.SetData(model.mIndiceGroups[i].data())
					.SetSize(model.mIndiceGroups[i].size() * sizeof(uint32))
					.SetIndexCount(model.mIndiceGroups[i].size())
					.Build();
			}



			mDescriptor = RHIDescriptor::Builder()
				.PushShader<PreDepthVertexShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("cameraUbo", mCameraBuffer);
			mDescriptor->SetValue("objectUbo", mSceneEntity->mObjectBuffer);
			mDescriptorMask = mDescriptor->AllocateDescriptor();

			

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
				.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "depthImage");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_NONE)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "renderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void PreDepthProcess::OnTick(float deltaTime)
{
	CameraUbo cameraUbo;
	cameraUbo.proj = mGameScene->GetCamera().GetProj();
	cameraUbo.view = mGameScene->GetCamera().GetView();
	cameraUbo.cameraPos = mGameScene->GetCamera().GetPosition();

	ObjectUbo objectUbo;
	objectUbo.model = mSceneEntity->GetTransform().GetWorldMatrix();

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, objectUbo, cameraUbo](VkCommandBuffer commandBuffer)
		{
			CameraUbo cameraUboData = cameraUbo;

			mCameraBuffer->WriteMemory(&cameraUboData);


			ObjectUbo objectUboData = objectUbo;

			mSceneEntity->mObjectBuffer->WriteMemory(&objectUboData);




			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (int32 i = 0; i < mObjectSize; ++i)
			{
				/*if (i >= mLocalScene->mObjectIndex)
				{
					break;
				}*/

				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = mSceneEntity->mVertexBuffers[i]->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = mSceneEntity->mIndexBuffers[i]->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdDrawIndexed(commandBuffer, mSceneEntity->mIndexBuffers[i]->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void PreDepthProcess::OnResize()
{
	

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
		.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "depthImage");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
